import os
from PIL import Image
import numpy as np
import torch
from torch.utils.data import Dataset
from preprocess.transforms import ToTensor,Compose,RandomRotate,Center_Crop,RandomAffine,ColorJitter,RandomFlip
from torchvision import transforms


# https://pytorch.org/vision/master/_modules/torchvision/transforms/functional.html
class TrainDataset(Dataset):
    def __init__(self, image_path_list):
        # generate image path list

        self.image_path_list = image_path_list

        self.transforms =  Compose([
            ToTensor(),
            RandomFlip(),
            RandomRotate(180),
            RandomAffine(degrees=10, translate=(0.1,0.1)),
            Center_Crop(160),
            # ColorJitter(brightness=0.3,contrast=0.3,saturation=0.3)


        ])
        self .img_transform = transforms.Compose(
            [transforms.RandomErasing(scale=(0.02, 0.33)),]
        )
    def __getitem__(self, index):
        image_path,mask_path = self.image_path_list[index]
        image = Image.open(image_path)
        mask = Image.open(mask_path)

        # get label from dirname
        label = int(image_path.split("/")[-3])
        label = torch.tensor(label)
        mask = np.asarray(mask)
        mask[0] = Image.fromarray(np.multiply(np.asarray(image)[0],mask[0]))
        mask[1] = Image.fromarray(np.multiply(np.asarray(image)[1],mask[1]))
        mask[2] = Image.fromarray(np.multiply(np.asarray(image)[2],mask[2]))
        image,mask = self.transforms(image,mask)
        # image = self.img_transform(image)
        input0 = torch.stack([image[0],mask[0]],0)
        input1 = torch.stack([image[1],mask[1]],0)
        input2= torch.stack([image[2],mask[2]],0)
        input = torch.stack([input0,input1,input2],0)
        return input,label

    def __len__(self):
        return len(self.image_path_list)
